import pandas as pd
import arrow


# df columns=date id q ; q has + and _

def inven_everyday(df, begin, end):
    df1 = df.groupby(by=['id', 'date']).sum()
    # 可以把复合index重新拆分成columns

    trades = df1.reset_index()
    trades['id'] = trades['id']
    trades['date'] = trades['date']

    # create a range of dates for the merged dataframe pd.datetime.today()
    if not begin:
        begin = df['date'].min()
    if not end:
        end = df['date'].max()

    index_of_dates = pd.date_range(begin, end).to_frame().reset_index(drop=True).rename(columns={0: 'date'}).astype(str)
    # create a merged dataframe with columns date / stock / stock_Tr.
    merged = pd.merge(index_of_dates, trades, how='left', on='date')
    # print(merged)
    # create a pivottable showing the shares_TR of each stock for each date
    shares_tr = merged.pivot(index='date', columns='id', values='q')
    shares_tr = shares_tr.dropna(axis=1, how='all').fillna(0)
    cumShares = shares_tr.cumsum()
    cumShares.index = cumShares.index.astype(str)
    return cumShares


# 必须有goodsId，price date
def continous_price(df):
    # e = df['date'].max()
    e = arrow.now("Asia/Shanghai").format()[:10]
    s = df['date'].min()
    per = pd.date_range(s, e)
    per_s = per.strftime('%Y-%m-%d').tolist()
    df0 = pd.DataFrame({"date": per_s})
    ids = set(df['goodsId'].tolist())
    res = []
    for id in ids:
        df1 = df[df['goodsId'] == id]
        if len(df1) == 0:
            continue
        df2 = pd.merge(df0, df1, on="date", how="left")
        df2.sort_values(by=['date'], inplace=True)
        df2["price"] = df2['price'].fillna(method="ffill")
        df2["goodsId"] = df2['goodsId'].fillna(method="ffill")
        df3 = df2.dropna()
        df3.loc[:, 'price'] = df3['price'].astype(float)
        res.append(df3)
    df4 = pd.concat(res)
    return df4


def value_everyday(inven, price):
    # inven index is dates, columns is [id1,id2]
    inven = inven.stack()
    # now inven index is (date,id)
    df = pd.DataFrame({
        "date": [i[0] for i in inven.index],
        "goodsId": [i[1] for i in inven.index],
        "weight": inven.values.tolist()
    })

    df1 = pd.merge(df, price, on=['date', "goodsId"], how="left")
    df1.sort_values(by=['date'], inplace=True)
    df1.dropna(subset=['price', 'weight'], inplace=True)
    df1['value'] = df1['weight'] * df1['price']
    return df1


def continous_price_pivot(df):
    begin = df['date'].min()
    # end = df['date'].max()
    end = arrow.now("Asia/Shanghai").format()[:10]
    index_of_dates = pd.date_range(begin, end).to_frame().reset_index(drop=True).rename(columns={0: 'date'}).astype(str)
    df1 = df.groupby(by=['goodsId', 'date']).first()
    # 可以把复合index重新拆分成columns
    df2 = df1.reset_index()

    merged = pd.merge(index_of_dates, df2, how='left', on='date')
    df3 = merged.pivot(index='date', columns='goodsId', values='price')
    df3.dropna(how="all", axis=1, inplace=True)
    df3 = df3.fillna(method="ffill")
    df4 = df3.stack()
    df5 = pd.DataFrame({
        "date": [i[0] for i in df4.index],
        "goodsId": [i[1] for i in df4.index],
        "price": df4.values
    })
    return df5


if __name__ == '__main__':
    dates = ['2021-12-01', '2021-12-01', '2021-12-03', '2021-12-03', ]
    stock = ['A', 'B', 'A', 'B']
    q = [5, 10, 3, 1]
    df = pd.DataFrame({"date": dates, "goodsId": stock, "price": q})
    print(continous_price(df))
    print(continous_price_pivot(df))
